;;;Iterative methods: Jacobi and Gauss-Siedel methods

(load "/home/cl/gits/scheme/numerical-analysis/matrix-vector-operations.scm")
(library-directories '("." "/home/cl/gits/libs-for-chezscheme"))
(import (range)
        (parallelization-in-vector))

;;;;;;----------------------- Basic iteration method: give x -> new x -----------------------
;;;---- Jacobi method
;;note: A_ii != 0
;;n=dim(x)
(define (nextstate-Jacobi A b x n)
  (vector-map/parallel
   (lambda (i)
     (/ (- (vector-ref b i)
           (apply + (map (lambda (j)
                           (if (= j i)
                               0
                               (* (matrix-ref A i j)
                                  (vector-ref x j))))
                         (range n))))
        (matrix-ref A i i)))
   (vector-range n)))

;;;---- Gauss-Siedel method
(define (nextstate-GaussSiedel A b x n)
  (let ([newx (vector-copy x)])
    (vector-for-each!/parallel
     (lambda (i)
       (vector-set! newx i
                    (/ (- (vector-ref b i)
                          (do ((sum 0)
                               (j 0 (+ j 1)))
                              ((= j n) sum);;sum
                            (if (not (= j i))
                                (set! sum (+ sum
                                             (* (matrix-ref A i j)
                                                (vector-ref newx j);newx not x
                                                ))))))
                       (matrix-ref A i i))))
     (vector-range n))
    newx))


;;;---------------- Iterative methods: Orginal version ----------------------

(define (Iterative-method method A b x0 TOL MaxSteps)
  ;;method is a symbol
  (define nextstate (cond ((eq? method 'Jacobi)
                           nextstate-Jacobi)
                          ((eq? method 'GaussSiedel)
                           nextstate-GaussSiedel)
                          (else 
                           (error 'nextstate "No such method: " method))))
  (define n (vector-length b))
  ;; 计算前检查
  (assert (not (equal? b (make-vector n 0))));;b 不能为 0, 否则0就是解
  (assert (= n (vector-length A)))
  (assert (not (zero? (apply * (map (lambda (i)
                                    (matrix-ref A i i))
                                  (range n))))))
  ;;iteration
  (let iter ([x x0] [steps 0])
    (let* ([newx (nextstate A b x n)]
           [R (/ (vector-norm (vector-map - newx x))
                 (vector-norm newx))]
           [newsteps (+ steps 1)])
      (cond [(and (<= newsteps MaxSteps) (> R TOL))
             (iter newx newsteps)]
            [(and (<= newsteps MaxSteps) (<= R TOL))
             newx]
            [else (error #f "StepsNumber is Overflow, the method is failed: " method)])))
  )

(define (Jacobi A b x0 TOL MaxSteps)
  (Iterative-method 'Jacobi A b x0 TOL MaxSteps))
(define (GaussSiedel A b x0 TOL MaxSteps)
  (Iterative-method 'GaussSiedel A b x0 TOL MaxSteps))



;;;--------------------  Aitken acceleration --------------------------
;;note: 这里的加速还只是一次加速, 有超级加速器--见sicp 之 stream
(define (Iterative-method-with-Aitken-acceleration method A b x0 TOL MaxSteps)
  (define nextstate (cond ((eq? method 'Jacobi)
                           nextstate-Jacobi)
                          ((eq? method 'GaussSiedel)
                           nextstate-GaussSiedel)
                          (else 
                           (error "No such method: " method))))
  (define n (vector-length b))
  ;;计算前检查
  (assert (= n (vector-length A)))
  (assert (not (equal? b (make-vector n 0.0))))
  (assert (not (zero? (apply * (map (lambda (i) (matrix-ref A i i))
                                    (range n))))))
  ;; acceleration: (x y z) => newx
  (let iter ([x x0] [steps 0])
    (let* ([y (nextstate A b x n)]
           [z (nextstate A b y n)]
           [denomi (vector-map (lambda (ex ey ez) (+ ex ez (* -2 ey)))
                               x y z)];;分母
           [newx (vector-map (lambda (ex ey ez ede)
                               (if (= ede 0) ez ;;分母为0, 无法加速
                                   (- ez (/ (expt (- ez ey) 2) ede))))
                             x y z denomi)]
           [R (/ (vector-norm (vector-map - newx z))
                 (vector-norm newx))]
           [newsteps (+ steps 3)]);;since y,z,newx
      (cond [(and (<= newsteps MaxSteps) (> R TOL))
             (iter newx newsteps)]
            [(and (<= newsteps MaxSteps) (<= R TOL))
             newx]
            [else
             (error #f "StepsNumber is Overflow, failed Aikten method: " method)]))))

(define (Jacobi-Aitken A b x0 TOL MaxSteps)
  (Iterative-method-with-Aitken-acceleration 'Jacobi A b x0 TOL MaxSteps))
(define (GaussSiedel-Aitken A b x0 TOL MaxSteps)
  (Iterative-method-with-Aitken-acceleration 'GaussSiedel A b x0 TOL MaxSteps))


#|
;;---test
(let* ([A (vector '#(10 -1 0)
                  '#(-1 10 -2)
                  '#(0 -4 10))]
       [b (vector 9 7 6)]
       [x0 (make-vector 3 1000.0)]
       [diff 0.0001])
  (newline)
  (display "precise solution is:            " )
  (display (make-vector 3 1))
  (newline)
  (display "Jacobi solution is:             " )
  (display (Jacobi A b x0 diff 100))
  ;;why does only Jacobi-Aikten method fail?
  ;;(newline)
  ;;(display "Jacobi-Aitken solution is:      " )
  ;;(display (Jacobi-Aitken A b x0 diff 100))
  (newline)
  (display "GaussSiedel solution is:        " )
  (display (GaussSiedel A b x0 diff 100))
  (newline)
  (display "GaussSiedel-Aitken solution is: " )
  (GaussSiedel-Aitken A b x0 diff 100)
  )

|#

